#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
调度系统性能测试脚本

用于测试调度系统API的性能指标，包括：
- 响应时间测试
- 并发性能测试
- 压力测试
- 内存使用测试
- 数据库性能测试

使用方法：
python performance_test.py
"""

import requests
import time
import threading
import statistics
import json
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import psutil
import os

# 配置
BASE_URL = "http://localhost:5000/api"
TEST_USERNAME = "admin"
TEST_PASSWORD = "admin123"

class PerformanceTester:
    def __init__(self, base_url=BASE_URL):
        self.base_url = base_url
        self.token = None
        self.results = {
            'response_times': [],
            'concurrent_results': [],
            'stress_results': [],
            'error_count': 0,
            'total_requests': 0
        }
        
    def login(self):
        """登录获取Token"""
        url = f"{self.base_url}/auth/login"
        data = {
            "username": TEST_USERNAME,
            "password": TEST_PASSWORD
        }
        
        try:
            response = requests.post(url, json=data)
            if response.status_code == 200:
                result = response.json()
                self.token = result['access_token']
                print(f"✓ 登录成功")
                return True
            else:
                print(f"✗ 登录失败: {response.text}")
                return False
        except Exception as e:
            print(f"✗ 登录异常: {e}")
            return False
    
    def make_request(self, method, endpoint, data=None, params=None):
        """发送HTTP请求并记录性能数据"""
        url = f"{self.base_url}{endpoint}"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.token}"
        }
        
        start_time = time.time()
        
        try:
            if method.upper() == "GET":
                response = requests.get(url, headers=headers, params=params, timeout=30)
            elif method.upper() == "POST":
                response = requests.post(url, headers=headers, json=data, timeout=30)
            elif method.upper() == "PUT":
                response = requests.put(url, headers=headers, json=data, timeout=30)
            elif method.upper() == "DELETE":
                response = requests.delete(url, headers=headers, timeout=30)
            else:
                raise ValueError(f"Unsupported method: {method}")
            
            end_time = time.time()
            response_time = (end_time - start_time) * 1000  # 转换为毫秒
            
            self.results['total_requests'] += 1
            self.results['response_times'].append(response_time)
            
            if response.status_code >= 400:
                self.results['error_count'] += 1
            
            return response, response_time
            
        except Exception as e:
            end_time = time.time()
            response_time = (end_time - start_time) * 1000
            self.results['error_count'] += 1
            self.results['total_requests'] += 1
            return None, response_time
    
    def test_response_times(self):
        """测试各API端点的响应时间"""
        print("\n=== 响应时间测试 ===")
        
        endpoints = [
            ("GET", "/auth/profile", None, None),
            ("GET", "/work_orders", None, {"page": 1, "per_page": 10}),
            ("GET", "/work_orders/stats", None, None),
            ("GET", "/users", None, {"page": 1, "per_page": 10}),
            ("GET", "/groups", None, None),
            ("GET", "/vehicles", None, None),
            ("GET", "/materials", None, None),
            ("GET", "/work_orders/fields", None, None)
        ]
        
        endpoint_results = {}
        
        for method, endpoint, data, params in endpoints:
            times = []
            print(f"\n测试 {method} {endpoint}")
            
            # 每个端点测试10次
            for i in range(10):
                response, response_time = self.make_request(method, endpoint, data, params)
                times.append(response_time)
                
                if response and response.status_code < 400:
                    status = "✓"
                else:
                    status = "✗"
                
                print(f"  {i+1}/10: {response_time:.2f}ms {status}")
                time.sleep(0.1)  # 避免过于频繁的请求
            
            # 计算统计数据
            avg_time = statistics.mean(times)
            min_time = min(times)
            max_time = max(times)
            median_time = statistics.median(times)
            
            endpoint_results[f"{method} {endpoint}"] = {
                'average': avg_time,
                'min': min_time,
                'max': max_time,
                'median': median_time,
                'times': times
            }
            
            print(f"  平均响应时间: {avg_time:.2f}ms")
            print(f"  最小/最大: {min_time:.2f}ms / {max_time:.2f}ms")
            print(f"  中位数: {median_time:.2f}ms")
        
        return endpoint_results
    
    def concurrent_request_worker(self, endpoint_info, request_count):
        """并发请求工作线程"""
        method, endpoint, data, params = endpoint_info
        results = []
        
        for _ in range(request_count):
            response, response_time = self.make_request(method, endpoint, data, params)
            results.append({
                'response_time': response_time,
                'success': response is not None and response.status_code < 400
            })
        
        return results
    
    def test_concurrent_performance(self):
        """测试并发性能"""
        print("\n=== 并发性能测试 ===")
        
        # 测试配置
        concurrent_levels = [1, 5, 10, 20]
        requests_per_thread = 5
        test_endpoint = ("GET", "/work_orders", None, {"page": 1, "per_page": 10})
        
        concurrent_results = {}
        
        for concurrent_count in concurrent_levels:
            print(f"\n测试并发数: {concurrent_count}")
            print(f"每线程请求数: {requests_per_thread}")
            print(f"总请求数: {concurrent_count * requests_per_thread}")
            
            start_time = time.time()
            
            # 使用线程池执行并发请求
            with ThreadPoolExecutor(max_workers=concurrent_count) as executor:
                futures = []
                for _ in range(concurrent_count):
                    future = executor.submit(self.concurrent_request_worker, test_endpoint, requests_per_thread)
                    futures.append(future)
                
                all_results = []
                for future in as_completed(futures):
                    thread_results = future.result()
                    all_results.extend(thread_results)
            
            end_time = time.time()
            total_time = end_time - start_time
            
            # 计算统计数据
            response_times = [r['response_time'] for r in all_results]
            success_count = sum(1 for r in all_results if r['success'])
            
            avg_response_time = statistics.mean(response_times)
            throughput = len(all_results) / total_time
            success_rate = (success_count / len(all_results)) * 100
            
            concurrent_results[concurrent_count] = {
                'total_time': total_time,
                'avg_response_time': avg_response_time,
                'throughput': throughput,
                'success_rate': success_rate,
                'total_requests': len(all_results),
                'successful_requests': success_count
            }
            
            print(f"  总耗时: {total_time:.2f}秒")
            print(f"  平均响应时间: {avg_response_time:.2f}ms")
            print(f"  吞吐量: {throughput:.2f} 请求/秒")
            print(f"  成功率: {success_rate:.1f}%")
        
        return concurrent_results
    
    def test_stress_performance(self):
        """压力测试"""
        print("\n=== 压力测试 ===")
        
        # 压力测试配置
        duration = 30  # 测试持续时间（秒）
        concurrent_users = 10
        test_endpoint = ("GET", "/work_orders", None, {"page": 1, "per_page": 5})
        
        print(f"测试持续时间: {duration}秒")
        print(f"并发用户数: {concurrent_users}")
        
        # 记录系统资源使用情况
        process = psutil.Process()
        initial_memory = process.memory_info().rss / 1024 / 1024  # MB
        
        start_time = time.time()
        end_time = start_time + duration
        
        results = []
        stop_flag = threading.Event()
        
        def stress_worker():
            """压力测试工作线程"""
            worker_results = []
            while not stop_flag.is_set():
                response, response_time = self.make_request(*test_endpoint)
                worker_results.append({
                    'timestamp': time.time(),
                    'response_time': response_time,
                    'success': response is not None and response.status_code < 400
                })
                time.sleep(0.1)  # 避免过于频繁
            return worker_results
        
        # 启动压力测试线程
        with ThreadPoolExecutor(max_workers=concurrent_users) as executor:
            futures = [executor.submit(stress_worker) for _ in range(concurrent_users)]
            
            # 等待测试时间结束
            time.sleep(duration)
            stop_flag.set()
            
            # 收集结果
            for future in as_completed(futures):
                worker_results = future.result()
                results.extend(worker_results)
        
        # 计算压力测试统计数据
        if results:
            response_times = [r['response_time'] for r in results]
            success_count = sum(1 for r in results if r['success'])
            
            avg_response_time = statistics.mean(response_times)
            max_response_time = max(response_times)
            min_response_time = min(response_times)
            throughput = len(results) / duration
            success_rate = (success_count / len(results)) * 100
            
            # 检查内存使用
            final_memory = process.memory_info().rss / 1024 / 1024  # MB
            memory_increase = final_memory - initial_memory
            
            stress_results = {
                'duration': duration,
                'total_requests': len(results),
                'successful_requests': success_count,
                'avg_response_time': avg_response_time,
                'min_response_time': min_response_time,
                'max_response_time': max_response_time,
                'throughput': throughput,
                'success_rate': success_rate,
                'initial_memory_mb': initial_memory,
                'final_memory_mb': final_memory,
                'memory_increase_mb': memory_increase
            }
            
            print(f"\n压力测试结果:")
            print(f"  总请求数: {len(results)}")
            print(f"  成功请求数: {success_count}")
            print(f"  平均响应时间: {avg_response_time:.2f}ms")
            print(f"  响应时间范围: {min_response_time:.2f}ms - {max_response_time:.2f}ms")
            print(f"  吞吐量: {throughput:.2f} 请求/秒")
            print(f"  成功率: {success_rate:.1f}%")
            print(f"  内存使用: {initial_memory:.1f}MB -> {final_memory:.1f}MB (增加 {memory_increase:.1f}MB)")
            
            return stress_results
        
        return None
    
    def test_database_performance(self):
        """测试数据库相关操作的性能"""
        print("\n=== 数据库性能测试 ===")
        
        # 测试大量数据查询
        print("\n1. 大量数据查询测试")
        large_query_times = []
        
        for page in range(1, 6):  # 测试前5页数据
            params = {"page": page, "per_page": 50}
            response, response_time = self.make_request("GET", "/work_orders", params=params)
            large_query_times.append(response_time)
            print(f"  第{page}页 (50条/页): {response_time:.2f}ms")
        
        avg_large_query_time = statistics.mean(large_query_times)
        print(f"  平均查询时间: {avg_large_query_time:.2f}ms")
        
        # 测试复杂查询
        print("\n2. 复杂查询测试")
        complex_queries = [
            {"search": "测试", "status": "pending", "priority": "high"},
            {"search": "维修", "status": "in_progress"},
            {"priority": "urgent", "page": 1, "per_page": 20}
        ]
        
        complex_query_times = []
        for i, params in enumerate(complex_queries, 1):
            response, response_time = self.make_request("GET", "/work_orders", params=params)
            complex_query_times.append(response_time)
            print(f"  复杂查询{i}: {response_time:.2f}ms")
        
        avg_complex_query_time = statistics.mean(complex_query_times)
        print(f"  平均复杂查询时间: {avg_complex_query_time:.2f}ms")
        
        # 测试统计查询
        print("\n3. 统计查询测试")
        stats_times = []
        for _ in range(5):
            response, response_time = self.make_request("GET", "/work_orders/stats")
            stats_times.append(response_time)
        
        avg_stats_time = statistics.mean(stats_times)
        print(f"  平均统计查询时间: {avg_stats_time:.2f}ms")
        
        return {
            'large_query_avg': avg_large_query_time,
            'complex_query_avg': avg_complex_query_time,
            'stats_query_avg': avg_stats_time
        }
    
    def generate_performance_report(self, response_results, concurrent_results, stress_results, db_results):
        """生成性能测试报告"""
        report = {
            "test_info": {
                "timestamp": datetime.now().isoformat(),
                "base_url": self.base_url,
                "total_requests": self.results['total_requests'],
                "error_count": self.results['error_count'],
                "error_rate": (self.results['error_count'] / self.results['total_requests'] * 100) if self.results['total_requests'] > 0 else 0
            },
            "response_time_tests": response_results,
            "concurrent_tests": concurrent_results,
            "stress_test": stress_results,
            "database_tests": db_results,
            "summary": {
                "overall_avg_response_time": statistics.mean(self.results['response_times']) if self.results['response_times'] else 0,
                "fastest_response": min(self.results['response_times']) if self.results['response_times'] else 0,
                "slowest_response": max(self.results['response_times']) if self.results['response_times'] else 0
            }
        }
        
        # 保存报告
        report_file = f"performance_report_{int(time.time())}.json"
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        print(f"\n性能测试报告已保存到: {report_file}")
        return report
    
    def run_all_tests(self):
        """运行所有性能测试"""
        print("开始调度系统性能测试...")
        print(f"测试目标: {self.base_url}")
        print("=" * 50)
        
        # 登录
        if not self.login():
            print("登录失败，停止测试")
            return
        
        try:
            # 运行各项性能测试
            response_results = self.test_response_times()
            concurrent_results = self.test_concurrent_performance()
            stress_results = self.test_stress_performance()
            db_results = self.test_database_performance()
            
            # 生成报告
            report = self.generate_performance_report(
                response_results, concurrent_results, stress_results, db_results
            )
            
            # 输出总结
            print("\n" + "=" * 50)
            print("性能测试总结:")
            print(f"总请求数: {self.results['total_requests']}")
            print(f"错误数: {self.results['error_count']}")
            print(f"错误率: {(self.results['error_count'] / self.results['total_requests'] * 100):.2f}%")
            
            if self.results['response_times']:
                print(f"平均响应时间: {statistics.mean(self.results['response_times']):.2f}ms")
                print(f"最快响应: {min(self.results['response_times']):.2f}ms")
                print(f"最慢响应: {max(self.results['response_times']):.2f}ms")
            
            # 性能评估
            avg_response = statistics.mean(self.results['response_times']) if self.results['response_times'] else 0
            error_rate = (self.results['error_count'] / self.results['total_requests'] * 100) if self.results['total_requests'] > 0 else 0
            
            print("\n性能评估:")
            if avg_response < 100:
                print("✓ 响应时间: 优秀 (<100ms)")
            elif avg_response < 500:
                print("✓ 响应时间: 良好 (<500ms)")
            elif avg_response < 1000:
                print("⚠ 响应时间: 一般 (<1000ms)")
            else:
                print("✗ 响应时间: 需要优化 (>1000ms)")
            
            if error_rate < 1:
                print("✓ 错误率: 优秀 (<1%)")
            elif error_rate < 5:
                print("✓ 错误率: 良好 (<5%)")
            else:
                print("✗ 错误率: 需要改进 (>5%)")
                
        except Exception as e:
            print(f"性能测试过程中发生错误: {e}")

def main():
    """主函数"""
    print("调度系统性能测试套件")
    print("=" * 30)
    
    tester = PerformanceTester()
    tester.run_all_tests()

if __name__ == "__main__":
    main()